import asyncio
from abc import abstractmethod, ABC
from typing import Union, List, Optional, AsyncIterable

import requests

from server.module_chat.chat.core.agent.tool.api.chat_result import ChatResult
from server.module_chat.chat.core.agent.tool.tool import ToolApi
from server.module_chat.chat.entity.vo.history import History


class ChatApi(ABC):
    """
    ChatApi
    """
    api_key: str

    api_base_url: str

    api_endpoint: str

    model: str

    temperature: float = 0.8

    max_tokens: Optional[int]

    top_p: float = 1.0

    presence_penalty: float = 1.0

    frequency_penalty: float = 1.0

    def __init__(self, tool_api: ToolApi, **kwargs):
        self.model = tool_api.model
        self.api_key = tool_api.api_key
        self.api_base_url = tool_api.api_base_url
        self.api_endpoint = tool_api.api_endpoint

    async def send_message(self,
                           query: str,
                           conversation_id: str,
                           history_len,
                           history: Union[int, List[History]],
                           stream: bool,
                           temperature: float,
                           max_tokens: Optional[int],
                           top_p: float
                           ) -> AsyncIterable[str]:
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

        try:
            data = self.build_data(query=query,
                                   conversation_id=conversation_id,
                                   history_len=history_len,
                                   history=history,
                                   stream=stream,
                                   temperature=temperature,
                                   max_tokens=max_tokens,
                                   top_p=top_p)
            response = requests.post(
                url=f"{self.api_base_url}{self.api_endpoint}",
                headers=headers,
                json=data,
                stream=True,
            )
            for chunk in response.iter_lines():
                result = self.build_result(chunk)
                if result is not None:
                    yield result.model_dump_json(by_alias=True)
                await asyncio.sleep(0.01)
        except Exception as e:
            print(e)
            yield ChatResult(**{
                "conversation_id": "",
                "message_id": "",
                "text": str(e)
            }).model_dump_json(by_alias=True)

    @abstractmethod
    def build_data(self) -> dict:
        pass

    @abstractmethod
    def build_result(self, chunk: bytes) -> ChatResult:
        pass
